from torch import nn
import torch
import torch.nn.functional as F


class MLP(nn.Module):  # 定义MLP模型，继承Module模块
    def __init__(self, num_classes=10):  # 定义初始化方法，0-9一共有10类，所以num_classes设置为10
        super(MLP, self).__init__()  # 调用MLP类的__init__方法
        self.flatten = nn.Flatten()  # 先将输入进来的图片像素压平
        self.linear1 = nn.Linear(784, 512)  # 因为图片大小是1*28*28=784，所以线性函数的输入是784，将输出设置为512
        self.linear2 = nn.Linear(512, 128)  # 上一层的输出即为这一层的输入，并将输出设置为128
        self.linear3 = nn.Linear(128, num_classes)  # 上一层的输出即为这一层的输入，输出设置为要分类的数量

    def forward(self, x):  # 定义前向传播函数，输入图像
        # x = x.view(-1, 28 * 28)
        x = self.flatten(x)  # 调用flatten，将图像压平
        x = F.relu(self.linear1(x))  # 将压平后的结果进行第一次线性运算，然后利用relu函数进行激活
        x = F.relu(self.linear2(x))  # 将上一层的输出结果进行第二次线性运算，然后利用relu函数进行激活
        x = F.softmax((self.linear3(x)), dim=1)  # 将上一层输出结果进行第三次线性运算，然后利用softmax函数进行激活，dim=1表示，激活后的结果每一行相加等于1
        return x  # 将结果输出


if __name__ == '__main__':  # 用于检验模型是否正确，输入一个符合要求的矩阵，看能否正确输出
    model = MLP()
    input = torch.ones((1, 1, 28, 28))
    output = model(input)
    print(output.shape)
